Back to Blog | LearnableLoopAI.com | Portfolio of Projects |

versioninfo() ## Julia version

Julia Version 1.10.5
Commit 6f3fdf7b362 (2024-08-27 14:19 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 12 × Intel(R) Core(TM) i7-8700B CPU @ 3.20GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, skylake)
Threads: 1 default, 0 interactive, 1 GC (on 12 virtual cores)

import Pkg
Pkg.activate(@__DIR__)
Pkg.add([
    "RxInfer", "BenchmarkTools", "Plots",
    "LaTeXStrings", "Distributions", "StableRNGs",
])
using RxInfer, BenchmarkTools, Random, LinearAlgebra, Plots, LaTeXStrings, Distributions, StableRNGs

  Activating project at `/workspaces/Kalman filtering and smoothing`
    Updating registry at `~/.julia/registries/General.toml`
   Resolving package versions...
  No Changes to `/workspaces/Kalman filtering and smoothing/Project.toml`
  No Changes to `/workspaces/Kalman filtering and smoothing/Manifest.toml`

Pkg.status()

Status `/workspaces/Kalman filtering and smoothing/Project.toml`
  [6e4b80f9] BenchmarkTools v1.8.0
  [31c24e10] Distributions v0.25.131
  [b964fa9f] LaTeXStrings v1.4.1
  [91a5bcdd] Plots v1.41.7
⌅ [86711068] RxInfer v3.10.1
  [860ef19b] StableRNGs v1.0.4
Info Packages marked with ⌅ have new versions available but compatibility constraints restrict them from upgrading. To see why use `status --outdated`

Kalman filtering and smoothing (Part 1)

This is an analysis of the RxInfer example at https://examples.rxinfer.com/categories/basic_examples/kalman_filtering_and_smoothing/
Some symbols have been changed
Some content has been added/modified
The preference is to make the math and code names align as much as possible
Spatial structure identifiers (e.g. vectors, matrices, cuboids)
- math
  - have a boldface e.g. \(\mathbf{x}\) or \(\boldsymbol{x}\) or \(\mathbf{A}\) or \(\mathbf{\mathbb{X}}\)
- code
  - have a underscore prefix e.g. _x or _A or _𝕏
- small case / cap case / blackboard case are generally used to discriminate between vectors, matrices, and cuboids
- start of alphabet / end of alphabet are generally used to discriminate between ‘system’ and ‘signal’
Time structure identifiers (i.e. time sequence/series)
- math
  - have a colon subscript e.g. \(x_:\) or \(x_{0:n}\) or \(x_{:n}\)
- code
  - have a ː (length mark) suffix e.g. xː
External (i.e. true, environment) states and parameters identifiers
- math
  - have a superscript * e.g. \(v^*\)
- code
  - have a superscript x e.g. vˣ
  - the ‘x’ in the code superscript is used to imitate the * superscript in the math

In the following set of examples the goal is to estimate hidden states of a Dynamical process where all hidden states are Gaussians.

We start our journey with a simple

1. multivariate Linear Gaussian State Space Model (LGSSM), which can be solved analytically. We then solve an
1. identification problem which does not have an analytical solution. Utimately, we show how RxInfer.jl can
1. deal with missing observations.

1 Gaussian Linear Dynamical System

LGSSM can be described with the following equations:

\[\begin{aligned} p(\mathbf{x}_t|\mathbf{x}_{t - 1}) & = \mathcal{N}(\mathbf{x}_t; \mathbf{A} \mathbf{x}_{t - 1}, \mathbf{\Sigma_x}),\\ p(\mathbf{y}_t|\mathbf{x}_t) & = \mathcal{N}(\mathbf{y}_t; \mathbf{C} \mathbf{x}_t, \mathbf{\Sigma_y}), \end{aligned}\]

where \(\mathbf{x}_t\) are hidden states, \(\mathbf{y}_t\) are noisy observations, \(\mathbf{A}\), \(\mathbf{C}\) are state transition and observation matrices, \(\mathbf{\Sigma_x}\) and \(\mathbf{\Sigma_y}\) are state transition noise and observation noise covariance matrices. For a more rigorous introduction to Linear Gaussian Dynamical systems we refer to Simo Sarkka, Bayesian Filtering and Smoothing book.

seed = 1234
rng = MersenneTwister(seed)

## We will model 2-dimensional observations with rotation matrix `Aˣ`
## To avoid clutter we also assume that matrices `Aˣ`, `Cˣ`, `Σˣₓ`, and `Σˣᵧ`
## are known and fixed for all time-steps
_xˣ₀ = [ 10.0, -10.0 ]
θˣ = π / 35
_Aˣ = [ cos(θˣ) -sin(θˣ); 
        sin(θˣ) cos(θˣ) ]
_Cˣ = diageye(2)
_Σˣₓ = diageye(2)
_Σˣᵧ = 25.0 .* diageye(2)
T = 300; ## number of observations

Next step, is to generate some synthetic data.

The Generative Process

State transition function (\(f_E\))

The state transition function provides the deterministic part of the state flow. The probabilistic part is provided by the system noise:

\[\mathbf{\dot{x}^*}_{t} = f_E(\mathbf{x^*}_{t-1}; \mathbf{A^*}) + \mathbf{\omega}^*_x = \mathbf{A^*} \mathbf{x^*}_{t-1} + \mathbf{\omega}^*_x\]

The *s indicate that the parameters and variables are the true unobserved values.

## state transition function
function fE(; _Aˣ, _xˣₜ₋₁)
    return _Aˣ*_xˣₜ₋₁
end
fE(_Aˣ=_Aˣ, _xˣₜ₋₁=_xˣ₀)

2-element Vector{Float64}:
 10.856136028986725
 -9.063349850918055

Observation generation function (\(g_E\))

The observation generation function provides the deterministic part of the observation. The probabilistic part is provided by the observation noise:

\[\mathbf{y}_{t} = g_E(\mathbf{x^*}_{t-1}; \mathbf{C^*}) + \mathbf{\omega}^*_y = \mathbf{C^*} \mathbf{x^*}_{t-1} + \mathbf{\omega}^*_y\]

The *s indicate that the parameters and variables are the true unobserved values.

## observation generation function
function gE(; _Cˣ, _xˣₜ)
    return _Cˣ*_xˣₜ
end
gE(; _Cˣ=_Cˣ, _xˣₜ=_xˣ₀)

2-element Vector{Float64}:
  10.0
 -10.0

## Data comes from either a simulation/lab (sim|lab) OR from the field (fld)
## Data are handled either in batches (batch) OR online as individual points (point)
## Batch data accumulates either
    ## along the depth/examples dimension/axis (into the screen/page), OR
        ## typical for supervised & unsupervised learning
    ## along the time dimension/axis (down the screen page)
        ## typical for sequential decision learning (reinforcement learning & active inference)
function sim_batch_data(rng, T, _Aˣ, _Cˣ, _Qˣ, _Rˣ) ## simulated batch data
    _xˣₜ₋₁ = _xˣ₀
    _fEː = Vector{Vector{Float64}}(undef, T)
    _xˣː = Vector{Vector{Float64}}(undef, T)
    _gEː = Vector{Vector{Float64}}(undef, T)
    _yˣː = Vector{Vector{Float64}}(undef, T)
    for t in 1:T
        _fEː[t] = fE(_Aˣ=_Aˣ, _xˣₜ₋₁=_xˣₜ₋₁)
        _xˣː[t] = rand(rng, MvNormal(_fEː[t], _Σˣₓ))

        _gEː[t] = gE(_Cˣ=_Cˣ, _xˣₜ=_xˣː[t])
        _yˣː[t] = rand(rng, MvNormal(_gEː[t], _Σˣᵧ))

        _xˣₜ₋₁ = _xˣː[t]
    end
    return _xˣː, _yˣː
end

sim_batch_data (generic function with 1 method)

_xˣː, _yː = sim_batch_data(rng, T, _Aˣ, _Cˣ, _Σˣₓ, _Σˣᵧ);

_xˣː

300-element Vector{Vector{Float64}}:
 [11.72348323093797, -9.965093666774871]
 [13.433953356799666, -6.662214695025911]
 [14.47940302752061, -5.948167903271925]
 [15.082367302413372, -2.773515689579621]
 [15.0190904154249, -1.0406633141659591]
 [16.61608054877271, -1.0869083912496964]
 [13.435259160640983, 0.332906617038909]
 [13.041178664848662, 0.9331868884552565]
 [13.20745099487252, 2.061986400670867]
 [13.866194467362085, 2.7240623329970597]
 [13.48602769414664, 2.8648329338168077]
 [11.813487642230692, 3.947720960793642]
 [10.402275516921657, 4.446976616165268]
 ⋮
 [-14.962542396537636, -16.708566420373018]
 [-12.885154178566715, -16.74934693640748]
 [-10.357686720735035, -17.59513576594144]
 [-8.900670351775474, -18.153037318140186]
 [-4.846198760964206, -18.008749910564465]
 [-2.4425949476217257, -17.86309514027969]
 [-0.1930330987239821, -16.86812800693013]
 [0.5984888183056473, -17.115109486088613]
 [1.999265873884699, -18.268734581119805]
 [3.247316037542661, -17.262857347138443]
 [4.137580852607579, -17.588597818510166]
 [6.107162996669496, -18.15648788117835]

_yː

300-element Vector{Vector{Float64}}:
 [9.2510894634168, -14.47966513610108]
 [16.098019767647358, -8.020891675199014]
 [11.676896336616727, -6.04462681172117]
 [10.94355014332877, -2.2230350579687435]
 [15.379672181181194, -8.557810600921213]
 [22.143569744302255, -6.620557958877578]
 [14.190137248801722, 4.1792979197118205]
 [6.642817613756749, 5.919772666742777]
 [13.917322286796729, 4.668349683011124]
 [10.042201942979187, -4.98306958673917]
 [10.583444968929523, 1.2876467033251646]
 [12.64267247980931, 1.9055307998303883]
 [4.268651886874915, 1.7383980173934095]
 ⋮
 [-16.310336258242465, -14.290271421281044]
 [-10.473101960412215, -13.506724712544282]
 [-1.7561084460229228, -15.011417600503886]
 [-9.272003783043301, -19.775893063100256]
 [-7.602820664459581, -23.251421895071957]
 [3.4344854947555037, -8.094842782373432]
 [1.1218576004967742, -19.10984945367171]
 [0.15328211552346116, -15.77054288956493]
 [6.376893388067707, -21.437054846029106]
 [12.340801468548573, -22.194772303997237]
 [-3.1281207214127154, -17.700570938696234]
 [5.013701061005852, -17.035797013986652]

Let’s plot our synthetic dataset. Lines represent our hidden states we want to estimate using noisy observations, which are represented as dots.

p = plot(title="Hidden states with noisy observations")
p = plot!(p, getindex.(_xˣː, 1), label="Hidden Signal " * L"x^*_1", color=:red)
p = scatter!(p, getindex.(_yː, 1), label=false, markersize=2, color=:red)

p = plot!(p, getindex.(_xˣː, 2), label="Hidden Signal " * L"x^*_2", color=:blue)
p = scatter!(p, getindex.(_yː, 2), label=false, markersize=2, color=:blue)
plot(p)

The Generative Model

To create a model we use GraphPPL package and @model macro:

@model function rotate_ssm(_yː, _x₀, _A, _C, _Σₓ, _Σᵧ)
    _x_prior ~ MvNormalMeanCovariance(mean(_x₀), cov(_x₀))
    _xₜ₋₁ = _x_prior
    for t in 1:length(_yː)
        _xː[t] ~ MvNormalMeanCovariance(_A*_xₜ₋₁, _Σₓ) ## `_x` is a sequence of hidden states
        _yː[t] ~ MvNormalMeanCovariance(_C*_xː[t], _Σᵧ) ## `_y` is a sequence of "clamped" observations
        _xₜ₋₁ = _xː[t]
    end
end

To run inference we also specify a prior for our first hidden state:

_xˣ₀ = MvNormalMeanCovariance(zeros(2), 100.0*diageye(2));

## For large number of observations you need to use limit_stack_depth = 100 option during model creation, e.g. 
## infer(..., options = (limit_stack_depth = 500, ))`
## We assume the Aˣ, Cˣ, Σˣₓ, Σˣᵧ are known, i.e. not hidden  
result = infer(
    model=       rotate_ssm(_x₀=_xˣ₀, _A=_Aˣ, _C=_Cˣ, _Σₓ=_Σˣₓ, _Σᵧ=_Σˣᵧ),
    data=        (_yː = _yː,),
    free_energy= true
);

result.posteriors

Dict{Symbol, Any} with 2 entries:
  :_xː      => MvNormalWeightedMeanPrecision{Float64, Vector{Float64}, Matrix{F…
  :_x_prior => MvNormalWeightedMeanPrecision(…

xmarginals  = result.posteriors[:_xː]
logevidence = -result.free_energy; ## given the analytical solution, free energy will be equal to the negative log evidence

p = plot(title="Estimated states from noisy observations")
p = plot!(p, getindex.(_xˣː, 1), label="Hidden Signal " * L"x^*_1", color=:red, linestyle=:dash)
p = plot!(p, getindex.(_xˣː, 2), label="Hidden Signal " * L"x^*_2", color=:blue, linestyle=:dash)

p = plot!(p, getindex.(mean.(xmarginals), 1), ribbon=getindex.(var.(xmarginals), 1) .|> sqrt, fillalpha=0.5, label="Estimated Signal " * L"x_1", color=:pink)
p = plot!(p, getindex.(mean.(xmarginals), 2), ribbon=getindex.(var.(xmarginals), 2) .|> sqrt, fillalpha=0.5, label="Estimated Signal " * L"x_2", color=:lightblue)
plot(p)

As we can see from our plot, estimated signal resembles closely to the real hidden states with small variance. We maybe also interested in the value for minus log evidence:

logevidence

1-element Vector{Float64}:
 -1882.2434870099432